{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "aJ_pmgxvGur9"
   },
   "source": [
    "# Assignment 2 - Variational Auto-Encoders\n",
    "## Deep Generative Modeling teaching material"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "mEneMITS2agU"
   },
   "source": [
    "#### Instructions on how to use this notebook:\n",
    "\n",
    "This notebook is hosted on ``Google Colab``. To be able to work on it, you have to create your own copy. Go to *File* and select *Save a copy in Drive*.\n",
    "\n",
    "You can also avoid using ``Colab`` entirely, and download the notebook to run it on your own machine. If you choose this, go to *File* and select *Download .ipynb*.\n",
    "\n",
    "The advantage of using **Colab** is that you can use a GPU. You can complete this assignment with a CPU, but it will take a bit longer. Furthermore, we encourage you to train using the GPU not only for faster training, but also to get experience with this setting. This includes moving models and tensors to the GPU and back. This experience is very valuable because for various models and large datasets (like large CNNs for ImageNet, or Transformer models trained on Wikipedia), training on GPU is the only feasible way.\n",
    "\n",
    "The default ``Colab`` runtime does not have a GPU. To change this, go to *Runtime - Change runtime type*, and select *GPU* as the hardware accelerator. The GPU that you get changes according to what resources are available at the time, and its memory can go from a 5GB, to around 18GB if you are lucky. If you are curious, you can run the following in a code cell to check:\n",
    "\n",
    "```sh\n",
    "!nvidia-smi\n",
    "```\n",
    "\n",
    "Note that despite the name, ``Google Colab`` does  not support collaborative work without issues. When two or more people edit the notebook concurrently, only one version will be saved. You can choose to do group programming with one person sharing the screen with the others, or make multiple copies of the notebook to work concurrently.\n",
    "\n",
    "**Submission:** Please bring your (partial) solution to instruction sessions. Then you can discuss it with intructors and your colleagues."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 416,
     "status": "ok",
     "timestamp": 1710958057214,
     "user": {
      "displayName": "Jakub Tomczak",
      "userId": "04705914613687700967"
     },
     "user_tz": -60
    },
    "id": "lBgoJIpdLI2Y",
    "outputId": "24f2cee7-1abc-4510-9f89-05f34762491b"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Wed Mar 20 18:07:36 2024       \n",
      "+---------------------------------------------------------------------------------------+\n",
      "| NVIDIA-SMI 535.104.05             Driver Version: 535.104.05   CUDA Version: 12.2     |\n",
      "|-----------------------------------------+----------------------+----------------------+\n",
      "| GPU  Name                 Persistence-M | Bus-Id        Disp.A | Volatile Uncorr. ECC |\n",
      "| Fan  Temp   Perf          Pwr:Usage/Cap |         Memory-Usage | GPU-Util  Compute M. |\n",
      "|                                         |                      |               MIG M. |\n",
      "|=========================================+======================+======================|\n",
      "|   0  Tesla T4                       Off | 00000000:00:04.0 Off |                    0 |\n",
      "| N/A   49C    P0              26W /  70W |    155MiB / 15360MiB |      0%      Default |\n",
      "|                                         |                      |                  N/A |\n",
      "+-----------------------------------------+----------------------+----------------------+\n",
      "                                                                                         \n",
      "+---------------------------------------------------------------------------------------+\n",
      "| Processes:                                                                            |\n",
      "|  GPU   GI   CI        PID   Type   Process name                            GPU Memory |\n",
      "|        ID   ID                                                             Usage      |\n",
      "|=======================================================================================|\n",
      "+---------------------------------------------------------------------------------------+\n"
     ]
    }
   ],
   "source": [
    "!nvidia-smi"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "tsdc7fDp40rQ"
   },
   "source": [
    "## Introduction\n",
    "\n",
    "In this assignment, we are going to implement a Variational Auto-Encoder (VAE). A VAE is a likelihood-based deep generative model that consists of a stochastic encoder (a variational posterior over latent variables), a stochastic decoder, and a marginal distribution over latent variables (a.k.a. a prior). The model was originally proposed in two concurrent papers:\n",
    "- [Kingma, D. P., & Welling, M. (2013). Auto-encoding variational bayes. arXiv preprint arXiv:1312.6114.](https://arxiv.org/abs/1312.6114)\n",
    "- [Rezende, Danilo Jimenez, Shakir Mohamed, and Daan Wierstra. \"Stochastic backpropagation and approximate inference in deep generative models.\" International conference on machine learning. PMLR, 2014.](https://proceedings.mlr.press/v32/rezende14.html)\n",
    "\n",
    "You can read more about VAEs in Chapter 4 of the following book:\n",
    "- [Tomczak, J.M., \"Deep Generative Modeling\", Springer, 2022](https://link.springer.com/book/10.1007/978-3-030-93158-2)\n",
    "\n",
    "In particular, the goals of this assignment are the following:\n",
    "\n",
    "- Understand how VAEs are formulated.\n",
    "- Implement components of VAEs using PyTorch.\n",
    "- Train and evaluate your VAE model on image data.\n",
    "\n",
    "This notebook is essential for preparing a report. Moreover, please remember to submit the final notebook together with the report (PDF)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "RvsuVNczG6pP"
   },
   "source": [
    "### Theory behind VAEs\n",
    "\n",
    "VAEs are latent variable models trained with variational inference. In general, the latent variable models define the following generative process:\n",
    "\\begin{align}\n",
    "1.\\ & \\mathbf{z} \\sim p_{\\lambda}(\\mathbf{z}) \\\\\n",
    "2.\\ & \\mathbf{x} \\sim p_{\\theta}(\\mathbf{x}|\\mathbf{z})\n",
    "\\end{align}\n",
    "\n",
    "In plain words, we assume that for observable data $\\mathbf{x}$, there are some latent (hidden) factors $\\mathbf{z}$. Then, the training objective is log-likelihood function of the following form:\n",
    "$$\n",
    "\\log p_{\\vartheta}(\\mathbf{x})=\\log \\int p_\\theta(\\mathbf{x} \\mid \\mathbf{z}) p_\\lambda(\\mathbf{z}) \\mathrm{d} \\mathbf{z} .\n",
    "$$\n",
    "\n",
    "The problem here is the intractability of the integral if the dependencies between random variables $\\mathbf{x}$ and $\\mathbf{z}$ are non-linear and/or the distributions are non-Gaussian.\n",
    "\n",
    "By introducing variational posteriors $q_{\\phi}(\\mathbf{z}|\\mathbf{x})$, we get the following lower bound (the Evidence Lower Bound, ELBO):\n",
    "$$\n",
    "\\log p_{\\vartheta}(\\mathbf{x}) \\geq \\mathbb{E}_{\\mathbf{z} \\sim q_\\phi(\\mathbf{z} \\mid \\mathbf{x})}\\left[\\log p_\\theta(\\mathbf{x} \\mid \\mathbf{z})\\right]-\\mathrm{KL}\\left(q_\\phi(\\mathbf{z} \\mid \\mathbf{x}) \\| p_\\lambda(\\mathbf{z})\\right) .\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "suzhlbWqxtD9"
   },
   "source": [
    "## IMPORTS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "BjxkigYLxpB7"
   },
   "outputs": [],
   "source": [
    "# DO NOT REMOVE!\n",
    "import os\n",
    "\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import torch\n",
    "\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "import torchvision\n",
    "from torchvision.datasets import MNIST"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 3,
     "status": "ok",
     "timestamp": 1710958057622,
     "user": {
      "displayName": "Jakub Tomczak",
      "userId": "04705914613687700967"
     },
     "user_tz": -60
    },
    "id": "Cm23hRm6CqGh",
    "outputId": "c21fc80a-1008-4cec-8a54-5348916940b4"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The available device is cuda\n"
     ]
    }
   ],
   "source": [
    "# DO NOT REMOVE!\n",
    "# Check if GPU is available and determine the device\n",
    "if torch.cuda.is_available():\n",
    "  device = 'cuda'\n",
    "else:\n",
    "  device = 'cpu'\n",
    "\n",
    "print(f'The available device is {device}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "81CxONpmMulC"
   },
   "outputs": [],
   "source": [
    "# DO NOT REMOVE! (unless you work locally)\n",
    "# mount drive: WE NEED IT FOR SAVING IMAGES! NECESSARY FOR GOOGLE COLAB!\n",
    "from google.colab import drive\n",
    "drive.mount('/content/gdrive')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "OoPb92zNM4UY"
   },
   "outputs": [],
   "source": [
    "# DO NOT REMOVE! (unless you work locally)\n",
    "# PLEASE CHANGE IT TO YOUR OWN GOOGLE DRIVE!\n",
    "images_dir = '/content/gdrive/My Drive/Colab/Results/'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "I3zs31tOyCmq"
   },
   "source": [
    "## Auxiliary functions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "DF0agzL7tDHK"
   },
   "source": [
    "Let us define some useful log-distributions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "LIBNVRNJtHSd"
   },
   "outputs": [],
   "source": [
    "# DO NOT REMOVE\n",
    "PI = torch.from_numpy(np.asarray(np.pi))\n",
    "EPS = 1.e-5\n",
    "\n",
    "def log_categorical(x, p, num_classes=256, reduction=None):\n",
    "    x_one_hot = F.one_hot(x.long(), num_classes=num_classes)\n",
    "    log_p = x_one_hot * torch.log(torch.clamp(p, EPS, 1. - EPS))\n",
    "    if reduction == 'mean':\n",
    "        return torch.mean(log_p, list(range(1, len(x.shape))))\n",
    "    elif reduction == 'sum':\n",
    "        return torch.sum(log_p, list(range(1, len(x.shape))))\n",
    "    else:\n",
    "        return log_p\n",
    "\n",
    "def log_bernoulli(x, p, reduction=None):\n",
    "    pp = torch.clamp(p, EPS, 1. - EPS)\n",
    "    log_p = x * torch.log(pp) + (1. - x) * torch.log(1. - pp)\n",
    "    if reduction == 'mean':\n",
    "        return torch.mean(log_p, list(range(1, len(x.shape))))\n",
    "    elif reduction == 'sum':\n",
    "        return torch.sum(log_p, list(range(1, len(x.shape))))\n",
    "    else:\n",
    "        return log_p\n",
    "\n",
    "def log_normal_diag(x, mu, log_var, reduction=None):\n",
    "    D = x.shape[1]\n",
    "    log_p = -0.5 * torch.log(2. * PI) - 0.5 * log_var - 0.5 * torch.exp(-log_var) * (x - mu)**2.\n",
    "    if reduction == 'mean':\n",
    "        return torch.mean(torch.sum(log_p, list(range(1, len(x.shape)))))\n",
    "    elif reduction == 'sum':\n",
    "        return torch.sum(torch.sum(log_p, list(range(1, len(x.shape)))))\n",
    "    else:\n",
    "        return log_p\n",
    "\n",
    "\n",
    "def log_standard_normal(x, reduction=None):\n",
    "    D = x.shape[1]\n",
    "    log_p = -0.5 * torch.log(2. * PI) - 0.5 * x**2.\n",
    "    if reduction == 'mean':\n",
    "        return torch.mean(log_p, list(range(1, len(x.shape))))\n",
    "    elif reduction == 'sum':\n",
    "        return torch.sum(log_p, list(range(1, len(x.shape))))\n",
    "    else:\n",
    "        return log_p"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Q2LLOs0kn7iw"
   },
   "source": [
    "## Implementing VAEs\n",
    "\n",
    "The goal of this assignment is to implement four classes:\n",
    "- `Encoder`: this class implements the encoder (variational posterior), $q_{\\phi}(\\mathbf{z}|\\mathbf{x})$.\n",
    "- `Decoder`: this class implements the decoded (the conditional likelihood), $p_{\\theta}(\\mathbf{x}|\\mathbf{z})$.\n",
    "- `Prior`: this class implements the marginal over latents (the prior), $p_{\\lambda}(\\mathbf{z})$.\n",
    "- `VAE`: this class combines all components."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "7cXhOwKAzW6Z"
   },
   "source": [
    "### Encoder\n",
    "We start with `Encoder`. Please remember that we assume the Gaussian variational posterior with a diagonal covariance matrix."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "MrwQXSuEoFfH"
   },
   "outputs": [],
   "source": [
    "# YOUR CODE GOES IN THIS CELL\n",
    "# NOTE: The class must containt the following functions:\n",
    "# (i) reparameterization\n",
    "# (ii) sample\n",
    "# Moreover, forward must return the log-probability of variational posterior for given x, i.e., log q(z|x)\n",
    "\n",
    "class Encoder(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(Encoder, self).__init__()\n",
    "        #\n",
    "        # your code goes here\n",
    "        # NOTE: please use whatever arguments you want, but probably one of them is a neural net\n",
    "        #\n",
    "\n",
    "    @staticmethod\n",
    "    def reparameterization(mu, log_var):\n",
    "        #\n",
    "        # your code goes here\n",
    "        #\n",
    "        # return is the result of the reparameterization trick\n",
    "\n",
    "    def encode(self, x):\n",
    "        #\n",
    "        # your code goes here\n",
    "        #\n",
    "        # return is means and log-variances (or variances)\n",
    "\n",
    "    def sample(self, x):\n",
    "        #\n",
    "        # your code goes here\n",
    "        #\n",
    "        # return is a sample from a Gaussian variational posterior\n",
    "\n",
    "\n",
    "    def log_prob(self, x):\n",
    "        #\n",
    "        # your code goes here\n",
    "        #\n",
    "        # return is the log-probability of a Gaussian with diagonal covariance matrix\n",
    "\n",
    "\n",
    "    def forward(self, x):\n",
    "        #\n",
    "        # your code goes here\n",
    "        #\n",
    "        # return is the log-probability"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "nhNTy5mn0XDT"
   },
   "source": [
    "### Decoder\n",
    "\n",
    "The decoder is the conditional likelihood, i.e., $p(x|z)$. Please remember that we must decide on the form of the distribution (e.g., Bernoulli, Gaussian, Categorical)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "9vTmKHwrpUVa"
   },
   "outputs": [],
   "source": [
    "# YOUR CODE GOES IN THIS CELL\n",
    "# NOTE: The class must containt the following function:\n",
    "# (i) sample\n",
    "# Moreover, forward must return the log-probability of the conditional likelihood function for given z, i.e., log p(x|z)\n",
    "\n",
    "class Decoder(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(Decoder, self).__init__()\n",
    "        #\n",
    "        # your code goes here\n",
    "        # NOTE: please use whatever arguments you want, but probably one of them is a neural net\n",
    "        #\n",
    "\n",
    "    def decode(self, z):\n",
    "        #\n",
    "        # your code goes here\n",
    "        #\n",
    "        # return is parameters of the conditional distribution (decoder)\n",
    "\n",
    "    def sample(self, z):\n",
    "        #\n",
    "        # your code goes here\n",
    "        #\n",
    "        # return is a sample\n",
    "\n",
    "        return x_new\n",
    "\n",
    "    def log_prob(self, x, z):\n",
    "        #\n",
    "        # your code goes here\n",
    "        #\n",
    "        # return is the log-probability\n",
    "\n",
    "    def forward(self, z):\n",
    "        #\n",
    "        # your code goes here\n",
    "        #\n",
    "        # return is the log-probability"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "CLIEwIiw00op"
   },
   "source": [
    "### Prior\n",
    "\n",
    "The prior is the marginal distribution over latent variables, i.e., $p(z)$. It plays a crucial role in the generative process and also in synthesizing images of a better quality.\n",
    "\n",
    "In this assignment, you are asked to implement a prior that is learnable (e.g., parameterized by a neural network). If you decide to implement the standard Gaussian prior only, then please be aware that you will not get any points.\n",
    "\n",
    "For the learnable prior you can choose the **Mixture of Gaussians**."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "xQIvee5Cp69V"
   },
   "outputs": [],
   "source": [
    "# YOUR CODE GOES IN THIS CELL\n",
    "# NOTES:\n",
    "# (i) Implementing the standard Gaussian prior does not give you any points!\n",
    "# (ii) The function \"sample\" must be implemented.\n",
    "# (iii) The function \"forward\" must return the log-probability, i.e., log p(z)\n",
    "\n",
    "class Prior(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(Prior, self).__init__()\n",
    "        #\n",
    "        # your code goes here\n",
    "        # NOTE: please use whatever arguments you want, but probably one of them is a neural network\n",
    "        #\n",
    "\n",
    "    def sample(self, batch_size):\n",
    "        #\n",
    "        # your code goes here\n",
    "        #\n",
    "        # return is a sample\n",
    "\n",
    "    def log_prob(self, z):\n",
    "        #\n",
    "        # your code goes here\n",
    "        #\n",
    "        # return is the log-probability"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "aOM4QM9I_62d"
   },
   "source": [
    "### Complete VAE\n",
    "\n",
    "The last class is `VAE` tha combines all components. Please remember that this class must implement the **Negative ELBO** in `forward`, as well as `sample` (*hint*: it is a composition of `sample` functions from the prior and the decoder)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "OQpf-BeSqA6V"
   },
   "outputs": [],
   "source": [
    "# YOUR CODE GOES HERE\n",
    "# This class combines Encoder, Decoder and Prior.\n",
    "# NOTES:\n",
    "# (i) The function \"sample\" must be implemented.\n",
    "# (ii) The function \"forward\" must return the negative ELBO. Please remember to add an argument \"reduction\" that is either \"mean\" or \"sum\".\n",
    "class VAE(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(VAE, self).__init__()\n",
    "        #\n",
    "        # your code goes here\n",
    "        # NOTE: please use whatever arguments you want, but probably you may init here your prior, encoder and decoder\n",
    "        #\n",
    "\n",
    "    def sample(self, batch_size=64):\n",
    "        #\n",
    "        # your code goes here\n",
    "        #\n",
    "        # return is a sample\n",
    "\n",
    "    def forward(self, x, reduction='mean'):\n",
    "        #\n",
    "        # your code goes here\n",
    "        # NOTE: you must implement the Negative ELBO here! Don't forget to allow \"sum\" and \"mean\" reduction\n",
    "        #\n",
    "        # return is the Negative ELBO"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "hLhgze7DA4yx"
   },
   "source": [
    "### Evaluation and training functions\n",
    "\n",
    "**Please DO NOT remove or modify them.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "I9Dr3a6lqJ0W"
   },
   "outputs": [],
   "source": [
    "# DO NOT REMOVE\n",
    "def evaluation(test_loader, name=None, model_best=None, epoch=None):\n",
    "    # EVALUATION\n",
    "    if model_best is None:\n",
    "        # load best performing model\n",
    "        model_best = torch.load(name + '.model')\n",
    "\n",
    "    model_best.eval()\n",
    "    loss = 0.\n",
    "    N = 0.\n",
    "    for indx_batch, (test_batch, _) in enumerate(test_loader):\n",
    "        test_batch = test_batch.to(device)\n",
    "        loss_t = model_best.forward(test_batch, reduction='sum')\n",
    "        loss = loss + loss_t.item()\n",
    "        N = N + test_batch.shape[0]\n",
    "    loss = loss / N\n",
    "\n",
    "    if epoch is None:\n",
    "        print(f'FINAL LOSS: nll={loss}')\n",
    "    else:\n",
    "        print(f'Epoch: {epoch}, val nll={loss}')\n",
    "\n",
    "    return loss\n",
    "\n",
    "\n",
    "def samples_real(name, test_loader, shape=(28,28)):\n",
    "    # real images-------\n",
    "    num_x = 4\n",
    "    num_y = 4\n",
    "    x, _ = next(iter(test_loader))\n",
    "    x = x.to('cpu').detach().numpy()\n",
    "\n",
    "    fig, ax = plt.subplots(num_x, num_y)\n",
    "    for i, ax in enumerate(ax.flatten()):\n",
    "        plottable_image = np.reshape(x[i], shape)\n",
    "        ax.imshow(plottable_image, cmap='gray')\n",
    "        ax.axis('off')\n",
    "\n",
    "    plt.savefig(name+'_real_images.pdf', bbox_inches='tight')\n",
    "    plt.close()\n",
    "\n",
    "\n",
    "def samples_generated(name, data_loader, shape=(28,28), extra_name=''):\n",
    "    x, _ = next(iter(data_loader))\n",
    "    x = x.to('cpu').detach().numpy()\n",
    "\n",
    "    # generations-------\n",
    "    model_best = torch.load(name + '.model')\n",
    "    model_best.eval()\n",
    "\n",
    "    num_x = 4\n",
    "    num_y = 4\n",
    "    x = model_best.sample(num_x * num_y)\n",
    "    x = x.to('cpu').detach().numpy()\n",
    "\n",
    "    fig, ax = plt.subplots(num_x, num_y)\n",
    "    for i, ax in enumerate(ax.flatten()):\n",
    "        plottable_image = np.reshape(x[i], shape)\n",
    "        ax.imshow(plottable_image, cmap='gray')\n",
    "        ax.axis('off')\n",
    "\n",
    "    plt.savefig(name + '_generated_images' + extra_name + '.pdf', bbox_inches='tight')\n",
    "    plt.close()\n",
    "\n",
    "\n",
    "def plot_curve(name, nll_val):\n",
    "    plt.plot(np.arange(len(nll_val)), nll_val, linewidth='3')\n",
    "    plt.xlabel('epochs')\n",
    "    plt.ylabel('nll')\n",
    "    plt.savefig(name + '_nll_val_curve.pdf', bbox_inches='tight')\n",
    "    plt.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "9ABgMeG0qFAP"
   },
   "outputs": [],
   "source": [
    "# DO NOT REMOVE\n",
    "def training(name, max_patience, num_epochs, model, optimizer, training_loader, val_loader, shape=(28,28)):\n",
    "    nll_val = []\n",
    "    best_nll = 1000.\n",
    "    patience = 0\n",
    "\n",
    "    # Main loop\n",
    "    for e in range(num_epochs):\n",
    "        # TRAINING\n",
    "        model.train()\n",
    "        for indx_batch, (batch, _) in enumerate(training_loader):\n",
    "            batch = batch.to(device)\n",
    "            loss = model.forward(batch, reduction='mean')\n",
    "\n",
    "            optimizer.zero_grad()\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "\n",
    "        # Validation\n",
    "        loss_val = evaluation(val_loader, model_best=model, epoch=e)\n",
    "        nll_val.append(loss_val)  # save for plotting\n",
    "\n",
    "        if e == 0:\n",
    "            print('saved!')\n",
    "            torch.save(model, name + '.model')\n",
    "            best_nll = loss_val\n",
    "        else:\n",
    "            if loss_val < best_nll:\n",
    "                print('saved!')\n",
    "                torch.save(model, name + '.model')\n",
    "                best_nll = loss_val\n",
    "                patience = 0\n",
    "\n",
    "                samples_generated(name, val_loader, shape=shape, extra_name=\"_epoch_\" + str(e))\n",
    "            else:\n",
    "                patience = patience + 1\n",
    "\n",
    "        if patience > max_patience:\n",
    "            break\n",
    "\n",
    "    nll_val = np.asarray(nll_val)\n",
    "\n",
    "    return nll_val"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "kWr8N2u2qNTu"
   },
   "source": [
    "### Setup\n",
    "\n",
    "**NOTE: *Please comment your code! Especially if you introduce any new variables (e.g., hyperparameters).***\n",
    "\n",
    "In the following cells, we define `transforms` for the dataset. Next, we initialize the data, a directory for results and some fixed hyperparameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "bFTE5jtYpxDV"
   },
   "outputs": [],
   "source": [
    "# DO NOT REMOVE\n",
    "transforms_train = torchvision.transforms.Compose([torchvision.transforms.ToTensor(),\n",
    "                                             torchvision.transforms.Lambda(lambda x: torch.bernoulli(x)),\n",
    "                                             torchvision.transforms.Lambda(lambda x: torch.flatten(x,0)),\n",
    "                                             ],\n",
    "                                            )\n",
    "\n",
    "transforms_test = torchvision.transforms.Compose([torchvision.transforms.ToTensor(),\n",
    "                                             torchvision.transforms.Lambda(lambda x: torch.bernoulli(x)),\n",
    "                                             torchvision.transforms.Lambda(lambda x: torch.flatten(x,0)),\n",
    "                                             ],\n",
    "                                            )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "8SDcOBbGCM8z"
   },
   "source": [
    "Please do not modify the code in the next cell."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "UXHitzrYqNhY"
   },
   "outputs": [],
   "source": [
    "# DO NOT REMOVE\n",
    "#-dataset\n",
    "dataset = MNIST('/files/', train=True, download=True,\n",
    "                      transform=transforms_train\n",
    "                )\n",
    "\n",
    "train_dataset, val_dataset = torch.utils.data.random_split(dataset, [50000, 10000], generator=torch.Generator().manual_seed(14))\n",
    "\n",
    "test_dataset = MNIST('/files/', train=False, download=True,\n",
    "                      transform=transforms_test\n",
    "                     )\n",
    "#-dataloaders\n",
    "batch_size = 32\n",
    "\n",
    "train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)\n",
    "val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False)\n",
    "test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)\n",
    "\n",
    "#-hyperparams (please do not modify them for the final report)\n",
    "num_epochs = 1000 # max. number of epochs\n",
    "max_patience = 20 # an early stopping is used, if training doesn't improve for longer than 20 epochs, it is stopped"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "yXuz71Aqbc2s"
   },
   "outputs": [],
   "source": [
    "# DO NOT REMOVE\n",
    "#-creating a dir for saving results\n",
    "name = 'vae'  # NOTE: if you run multiple experiments, you would overwrite results. Please modify this part if necessary.\n",
    "result_dir = images_dir + name + '/'\n",
    "if not(os.path.exists(result_dir)):\n",
    "  os.mkdir(result_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "kmKDXMI0B231"
   },
   "source": [
    "In the next cell, please initialize the model. Please remember about commenting your code!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "b73aaBDxqSYb"
   },
   "outputs": [],
   "source": [
    "# YOUR CODE COMES HERE:\n",
    "#\n",
    "# your code goes here\n",
    "#\n",
    "# use all necessary code to initialize your VAE\n",
    "\n",
    "model.to(device) # at the end, your model must be put on the available device"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "iC8AkWt4CURT"
   },
   "source": [
    "Please initialize the optimizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "a3nTSDe7ql08"
   },
   "outputs": [],
   "source": [
    "# PLEASE DEFINE YOUR OPTIMIZER\n",
    "#\n",
    "# your code goes here\n",
    "#\n",
    "# please do not forget to define hyperparameters!\n",
    "# please do it like this: optimizer = ..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "P5GrzUcHFweG"
   },
   "source": [
    "### Training and final evaluation\n",
    "\n",
    "In the following two cells, we run the training and the final evaluation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "VD7WuY6bqnBK"
   },
   "outputs": [],
   "source": [
    "# DO NOT REMOVE OR MODIFY\n",
    "# Training procedure\n",
    "nll_val = training(name=result_dir + name, max_patience=max_patience,\n",
    "                   num_epochs=num_epochs, model=model, optimizer=optimizer,\n",
    "                   training_loader=train_loader, val_loader=val_loader,\n",
    "                   shape=(28,28))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "JAuMt9_wquOI"
   },
   "outputs": [],
   "source": [
    "# DO NOT REMOVE OR MODIFY\n",
    "# Final evaluation\n",
    "test_loss = evaluation(name=result_dir + name, test_loader=test_loader)\n",
    "f = open(result_dir + name + '_test_loss.txt', \"w\")\n",
    "f.write(str(test_loss))\n",
    "f.close()\n",
    "\n",
    "samples_real(result_dir + name, test_loader)\n",
    "samples_generated(result_dir + name, test_loader, extra_name='_FINAL')\n",
    "\n",
    "plot_curve(result_dir + name, nll_val)"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "provenance": [
    {
     "file_id": "1C6CaZsSdXdc5fONOVrJy5iYKzr2DIg9j",
     "timestamp": 1607008727578
    }
   ],
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
